import asyncio
from urllib.parse import urlparse
import websockets
import json,os
from pprint import pprint


class WebSocketCDPClient:
    def __init__(self, tab,debug=False):
        self.tab = tab
        self.debug = debug
        self.counter = self._count_up()
        self.cdp_list = []
        self.response = []

        self.received_type=''
        self.data_packets=[]
        self.tab.run_cdp('Network.enable')
        self.filter_in_content=""

    def extract_between_braces(self,s):
        # 查找第一个 '{' 和最后一个 '}' 的位置
        first_brace_index = s.find('{')
        last_brace_index = s.rfind('}')
        
        # 确保 '{' 和 '}' 存在，并且 '{' 在 '}' 之前
        if first_brace_index != -1 and last_brace_index != -1 and first_brace_index < last_brace_index:
            # 提取第一个 '{' 和最后一个 '}' 之间的字符串
            return s[first_brace_index + 1 : last_brace_index]
        else:
            return ""  # 如果没有找到有效的 { 和 }，返回空字符串

    def save_dict_to_file(self,save_folder, save_name, content):
        content_str=content.get("body")
        if self.filter_in_content:
            if self.filter_in_content not in content_str:
                return
    
            
        try:
            json_obj=json.loads(content_str)
        except:
            content_str="{"+self.extract_between_braces(content_str)+"}"
            json_obj = json.loads(content_str)
        # content=json.dumps(content)
        # content=json.dumps(content, ensure_ascii=False)
        # 如果目标文件夹不存在，创建它
        if not os.path.exists(save_folder):
            os.makedirs(save_folder)
        
        # 组合文件路径
        file_path = os.path.join(save_folder, f"{save_name}.json")
        
        # 打开文件并写入内容
        # with open(file_path, 'w', encoding='utf-8') as f:
        #     f.write(content)

        # 保存为JSON文件，确保中文字符正确保存
        with open(file_path, 'w', encoding='utf-8') as f:
            json.dump(json_obj, f, ensure_ascii=False, indent=4)    
        
        print(f"数据内容已保存到 {file_path}")

    def _count_up(self):
        num = 1  # 从 1 开始
        while True:
            yield num  # 返回当前值
            num += 1  # 每次调用时递增

    def get_websocket_url_by_tab(self):
        """根据 Tab 获取 WebSocket URL"""
        url = self.tab._browser._chromium_options.address
        tab_id = self.tab.tab_id
        return f'ws://{url}/devtools/page/{tab_id}'

    def format_cdp(self, cdp_command):
        """格式化 CDP 命令，添加 id 和调整字段名称"""
        cdp_command['id'] = next(self.counter)  # 使用迭代器生成 ID
        if 'command' in cdp_command:
            cdp_command['method'] = cdp_command.pop('command')
        if 'parameters' in cdp_command:
            cdp_command['params'] = cdp_command.pop('parameters')
        return cdp_command

    def send_cdp(self, json_list):
        """格式化传入的 CDP 命令列表"""
        # 如果json_list不是列表，将其转换成列表
        if not isinstance(json_list, list):
            json_list = [json_list]
        self.cdp_list = [self.format_cdp(json) for json in json_list]
        # 连接并发送命令列表
        self.connect_and_send()
        return self.response
    
    def show(self,txt):
        if self.debug:
            print(txt)





    async def _connect_and_send(self, ws_url):
        """连接 WebSocket 并发送命令"""
        async with websockets.connect(ws_url) as websocket:
            self.show(f"WebSocket 连接建立 {self.tab.url}")
            self.response = []  # 清空响应列表
            for command in self.cdp_list:
                # 发送命令
                await websocket.send(json.dumps(command))
                self.show(f"Command sent: {command}")
                
                # 等待浏览器返回响应
                response = await websocket.recv()
                self.show(f"Response received: {response}")
                self.response.append(response)  # 将响应添加到列表中

    def connect_and_send(self):
        """连接并发送 CDP 命令"""
        ws_url = self.get_websocket_url_by_tab()
        # 运行异步任务
        asyncio.get_event_loop().run_until_complete(self._connect_and_send(ws_url))



    def set_listener(self, event_name, callback):
        """设置监听器，用于处理特定事件"""
        # self.event_handlers[event_name] = callback
        self.event_name=event_name
        
        self.tab.driver.set_callback(self.event_name, callback)

    def listen(self,event_name,callback):
        self.event_name=event_name        
        self.tab.driver.set_callback(self.event_name, callback)
    def run(self,cdp):
        self.tab.run_cdp(cdp)


    def _response(self,**kwargs):
        # self.data_queue.put(kwargs)
        print(kwargs)

    def get_response_body(self, request_id): 
        try:
            # 调用 CDP API 获取响应体数据
            response = self.tab.run_cdp("Network.getResponseBody", requestId=request_id)
        except Exception as e:
            # 捕获特定异常并记录日志或输出错误
            response = {"body": '{"info": "error", "message": "' + str(e) + '"}'}

        return response

    


        
    def received_call(self,**args):
        r=args['response']
        r_id=args['requestId']
        mime=r['mimeType']
        url=r['url']
        if self.received_type in mime:
            print(f'get data: {url} {mime} {r_id}')
            info=self.getResponseBody(r_id)
            print(info)
            print('\n'*2)
        else:
            pass 
    def received_data(self,type='png'):
        t=type
        def call(**args):               
            r=args['response']
            r_id=args['requestId']
            mime=r['mimeType']
            url=r['url']
            if t in mime:
                print(f'get data: {url} {mime} {r_id}')
                info=self.getResponseBody(r_id)
                print(info)
                print('\n'*2)
        return call        
    def received_data_pretty(self,type='png'):
        t=type
        self.count=1
        def call(**args):               
            r=args['response']
            r_id=args['requestId']
            mime=r['mimeType']
            url=r['url']
            _type=args['type']
            args['_type']=_type
            
            _url=r['url']
            _name:str=_url.split('/')[-1].split('?')[0]
            args['_name']=_name
            if   _type  in ['Preflight','Stylesheet']  :
                return
            if t in mime:       
                c=str(self.count).ljust(3)
                
                info=f'{c} {self.adjust(_name,20)}    {self.adjust(_type,10)} {r_id.ljust(10)} {mime.ljust(25)}  {self.adjust(_url,100)}'
                print(info)
                self.count+=1
                
                self.data_packets.append(args)     
    

        return call 
    def received_data_and_save(self,type='json'):
        t=type
        self.count=1        
        def call(**args):               
            r=args['response']
            r_id=args['requestId']
            mime=r['mimeType']            
            _type=args['type']
            args['_type']=_type
            
            _url=r['url']
            _name:str=_url.split('/')[-1].split('?')[0]
            
            sub_url=self.get_domain(_url)
            if   _type  in ['Preflight','Stylesheet']  :
                return
            if t in mime:
                response_body: dict = self.get_response_body(r_id)       
                self.save_dict_to_file(self.main_url_domain(), f'{sub_url}{_name}+{r_id}', response_body)         

        return call 

    def main_url_domain(self):
        return self.get_domain(self.tab.url)
    def get_domain(self,url):
        # 解析 URL
        parsed_url = urlparse(url)
        # 获取主机部分并提取主域名
        domain = parsed_url.netloc
        # 去除 www 部分，如果有的话
        domain = domain.replace('www.', '')
        return domain
    
    def adjust(self,_str, length):
        # 如果输入字符串长度小于目标长度，右侧填充空格
        if len(_str) < length:
            return _str.ljust(length)
        # 如果输入字符串长度大于目标长度，截取字符串
        else:
            str=_str[:8]+"..."+_str[-(length-8-3):]
            return str    




class CDP:
    """
    Chrome DevTools Protocol (CDP) 的封装类，提供了与浏览器交互的各种命令。
    """

    class Browser:
        """与浏览器相关的命令"""
        getVersion = 'Browser.getVersion'  # 获取浏览器版本信息
        getWindowForTarget = 'Browser.getWindowForTarget'  # 获取目标页面所在的窗口
        close = 'Browser.close'  # 关闭浏览器

    class Page:
        """与页面相关的命令"""
        navigate = 'Page.navigate'  # 导航到指定 URL
        reload = 'Page.reload'  # 刷新当前页面
        stopLoading = 'Page.stopLoading'  # 停止页面加载
        getNavigationHistory = 'Page.getNavigationHistory'  # 获取导航历史
        captureScreenshot = 'Page.captureScreenshot'  # 捕获页面截图
        startScreencast = 'Page.startScreencast'  # 开始屏幕录制
        enable = 'Page.enable'  # 启用页面监控

    class Network:
        """与网络相关的命令"""
        enable = 'Network.enable'  # 启用网络监控
        disable = 'Network.disable'  # 禁用网络监控
        setRequestInterception = 'Network.setRequestInterception'  # 设置请求拦截
        getResponseBody = 'Network.getResponseBody'  # 获取响应体
        responseReceived = 'Network.responseReceived'  # 响应收到事件
        dataReceived = 'Network.dataReceived'  # 数据接收事件
        requestWillBeSent = 'Network.requestWillBeSent'  # 请求即将发送事件

    class DOM:
        """与 DOM 相关的命令"""
        enable = 'DOM.enable'  # 启用 DOM 监控
        getDocument = 'DOM.getDocument'  # 获取页面 DOM
        querySelector = 'DOM.querySelector'  # 查询 DOM 元素
        setAttributeValue = 'DOM.setAttributeValue'  # 设置 DOM 元素属性
        removeNode = 'DOM.removeNode'  # 删除 DOM 元素节点
        attributeModified = 'DOM.attributeModified'  # DOM 属性修改事件

    class Log:
        """与日志相关的命令"""
        enable = 'Log.enable'  # 启用日志记录
        clear = 'Log.clear'  # 清除日志

    class Runtime:
        """执行 JavaScript 相关的命令"""
        evaluate = 'Runtime.evaluate'  # 执行 JS 代码并返回结果

    class Debugger:
        """调试相关的命令"""
        enable = 'Debugger.enable'  # 启用调试
        setBreakpoint = 'Debugger.setBreakpoint'  # 设置断点
        resume = 'Debugger.resume'  # 恢复代码执行

    class Emulation:
        """设备模拟相关的命令"""
        setDeviceMetricsOverride = 'Emulation.setDeviceMetricsOverride'  # 设置设备视口尺寸
        setUserAgentOverride = 'Emulation.setUserAgentOverride'  # 设置用户代理



# 使用示例：
# cdp = CDP()
# print(cdp.Page.navigate)  # 输出 'Page.navigate'
# print(cdp.Network.enable)  # 输出 'Network.enable'


  
    # def Network(self):
    #     responseReceived='Network.responseReceived'
    #     # return Network()
    # # Network='Network.responseReceived'


if __name__ == '__main__':
    
    aa=CDP.Page.reload
    print(aa)

    # bb=CDP()
    # print(bb.Page.reload)


# client = WebSocketCDPClient(tab)
# client.send_cdp_list(cdp_json_list)  # 传入一个 CDP 命令列表
# client.connect_and_send()  # 连接并发送命令
